import { AiImage } from "../core/AiImage.ts";
import { AiGlobalConfig } from "../AiGlobalConfig.ts";
import { SseClient } from "../core/client/sse/SseClient.ts";
import { AiClientListener } from "../core/AiClientListener.ts";
import { WebSocketClient } from "../core/client/ws/WebSocketClient.ts";
import { InnerEditor } from "../../core/AiEditor.ts";
import { HttpStreamSocketClient } from "../core/client/http/HttpSocketClient.ts";
import { AiClient } from "../core/AiClient.ts";
import { AiResponseListener } from "../core/AiResponseListener.ts";
import { CustomAiImageConfig } from "./CustomAiImageConfig.ts";
import { AiResponse } from "../core/AiResponse.ts";

export class CustomAiImage extends AiImage {
    constructor(aiImageName: string, editor: InnerEditor, globalConfig: AiGlobalConfig) {
        super(editor, globalConfig, aiImageName);
        this.aiImageConfig = {
            protocol: "http",
            ...globalConfig.images?.[aiImageName]
        } as CustomAiImageConfig
    }

    createAiClientUrl(): string {
        const config = this.aiImageConfig as CustomAiImageConfig;
        if (typeof config.url === "string") {
            return config.url as string;
        }
        return config.url?.();
    }

    createAiClient(url: string, listener: AiResponseListener): AiClient {
        const config = this.aiImageConfig as CustomAiImageConfig;
        const aiClientListener: AiClientListener = {
            onStart: listener.onStart,
            onStop: listener.onStop,
            onMessage: (bodyString) => {
                const config = this.aiImageConfig as CustomAiImageConfig;
                if (config.parseResponse) {
                    const aiResponse = config.parseResponse?.(bodyString);
                    if (aiResponse) listener.onMessage(aiResponse);
                } else {
                    const aiResponse = {
                        ...JSON.parse(bodyString)
                    } as AiResponse
                    listener.onMessage(aiResponse);
                }
            }
        };
        const clientConfig = {
            url,
            method: config.method || "post",
            headers: config.headers?.(),
        };
        return config.protocol === "sse" ? new SseClient(clientConfig, aiClientListener)
            : config.protocol === "http" ? new HttpStreamSocketClient(clientConfig, aiClientListener)
                : new WebSocketClient(url, aiClientListener)
    }

    wrapPayload(prompt: string) {
        const config = this.aiImageConfig as CustomAiImageConfig;
        if (config.wrapPayload) {
            return config.wrapPayload?.(prompt);
        } else {
            const payload = {
                "prompt": prompt,
                ...config.data,
            }
            return JSON.stringify(payload);
        }
    }
}